#include "WebSocketServer.h"
#include "http.h"
#include <sstream>
#include <cctype>
#include <algorithm>

// 声明全局防火墙配置
extern FirewallConfig firewall_config;
extern std::mutex firewall_mutex;
std::queue<LogMessage> log_queue;
std::mutex queue_mutex;
std::condition_variable queue_cv;
using namespace std;
void WebSocketServer::start() {
    running = true;
    server_thread = std::thread(&WebSocketServer::run, this);
}

void WebSocketServer::stop() {
    running = false;
    if (server_thread.joinable()) {
        server_thread.join();
    }
    
    std::lock_guard<std::mutex> lock(clients_mutex);
    for (int client : clients) {
        close(client);
    }
    clients.clear();
    
    if (server_fd != -1) {
        close(server_fd);
        server_fd = -1;
    }
}
std::string base64_encode(const unsigned char* input, int length) {
    BIO *bio, *b64;
    BUF_MEM *bufferPtr;

    b64 = BIO_new(BIO_f_base64());
    bio = BIO_new(BIO_s_mem());
    bio = BIO_push(b64, bio);

    BIO_set_flags(bio, BIO_FLAGS_BASE64_NO_NL);
    BIO_write(bio, input, length);
    BIO_flush(bio);
    BIO_get_mem_ptr(bio, &bufferPtr);

    std::string result(bufferPtr->data, bufferPtr->length);
    BIO_free_all(bio);

    return result;
}

void WebSocketServer::handleClient(int client_fd) {
    while (running) {
        char buffer[4096];
        ssize_t len = recv(client_fd, buffer, sizeof(buffer), 0);
        
        if (len <= 0) {
            if (len == 0) {
                std::cout << "客户端关闭连接" << std::endl;
            } else {
                perror("recv");
            }
            break;
        }
        
        // 解析WebSocket帧
        unsigned char opcode = buffer[0] & 0x0F;
        unsigned char mask_bit = buffer[1] >> 7;
        unsigned long payload_len = buffer[1] & 0x7F;
        
        size_t mask_offset = 2;
        if (payload_len == 126) {
            payload_len = (buffer[2] << 8) | buffer[3];
            mask_offset += 2;
        } else if (payload_len == 127) {
            // 不支持过长的消息
            continue;
        }
        
        unsigned char mask[4];
        if (mask_bit) {
            memcpy(mask, buffer + mask_offset, 4);
            mask_offset += 4;
        }
        
        // 处理有效载荷
        std::string payload;
        for (size_t i = 0; i < payload_len; i++) {
            char c = buffer[mask_offset + i];
            if (mask_bit) {
                c ^= mask[i % 4];
            }
            payload.push_back(c);
        }
        
        // 只处理文本帧
        if (opcode != 0x01) continue;
        
        std::cout << "收到配置更新: " << payload << std::endl;
        
        // 解析防火墙配置
        FirewallConfigMessage config_msg;
        std::istringstream iss(payload);
        std::string token;
        
        // 解析端口列表 (格式: ports=80,443,8080)
        if (payload.find("ports=") != std::string::npos) {
            size_t start = payload.find("ports=") + 6;
            size_t end = payload.find(';', start);
            std::string ports_str = (end == std::string::npos) ? 
                payload.substr(start) : payload.substr(start, end - start);
            
            std::istringstream port_stream(ports_str);
            while (std::getline(port_stream, token, ',')) {
                try {
                    int port = std::stoi(token);
                    config_msg.blocked_ports.push_back(port);
                } catch (...) {
                    // 忽略无效端口
                }
            }
        }
        
        // 解析SYN洪水阈值 (格式: syn_threshold=100)
        if (payload.find("syn_threshold=") != std::string::npos) {
            size_t start = payload.find("syn_threshold=") + 14;
            size_t end = payload.find(';', start);
            std::string val_str = (end == std::string::npos) ? 
                payload.substr(start) : payload.substr(start, end - start);
            try {
                config_msg.syn_flood_threshold = std::stoi(val_str);
            } catch (...) {
                // 使用默认值
            }
        }
        
        // 解析封禁时间 (格式: ban_time=60)
        if (payload.find("ban_time=") != std::string::npos) {
            size_t start = payload.find("ban_time=") + 9;
            size_t end = payload.find(';', start);
            std::string val_str = (end == std::string::npos) ? 
                payload.substr(start) : payload.substr(start, end - start);
            try {
                config_msg.syn_flood_ban_time = std::stoi(val_str);
            } catch (...) {
                // 使用默认值
            }
        }
        
        // 解析慢速攻击超时 (格式: slow_timeout=10)
        if (payload.find("slow_timeout=") != std::string::npos) {
            size_t start = payload.find("slow_timeout=") + 13;
            size_t end = payload.find(';', start);
            std::string val_str = (end == std::string::npos) ? 
                payload.substr(start) : payload.substr(start, end - start);
            try {
                config_msg.slow_dos_timeout = std::stoi(val_str);
            } catch (...) {
                // 使用默认值
            }
        }
        
        // 解析最小数据速率 (格式: min_rate=100)
        if (payload.find("min_rate=") != std::string::npos) {
            size_t start = payload.find("min_rate=") + 9;
            size_t end = payload.find(';', start);
            std::string val_str = (end == std::string::npos) ? 
                payload.substr(start) : payload.substr(start, end - start);
            try {
                config_msg.slow_dos_min_rate = std::stoi(val_str);
            } catch (...) {
                // 使用默认值
            }
        }
        
        // 更新防火墙配置
        {
            std::lock_guard<std::mutex> lock(firewall_mutex);
            firewall_config.blocked_ports.clear();
            for (int port : config_msg.blocked_ports) {
                firewall_config.blocked_ports.insert(port);
            }
            firewall_config.syn_flood_threshold = config_msg.syn_flood_threshold;
            firewall_config.syn_flood_ban_time = config_msg.syn_flood_ban_time;
            firewall_config.slow_dos_timeout = config_msg.slow_dos_timeout;
            firewall_config.slow_dos_min_rate = config_msg.slow_dos_min_rate;
        }
        
        std::cout << "防火墙配置已更新" << std::endl;
    }
    
    // 客户端断开连接
    {
        std::lock_guard<std::mutex> lock(clients_mutex);
        clients.erase(client_fd);
    }
    close(client_fd);
}

void WebSocketServer::run() {
    server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        perror("socket");
        return;
    }

    int opt = 1;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    sockaddr_in address{};
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(port);

    if (bind(server_fd, (struct sockaddr*)&address, sizeof(address)) < 0) {
        perror("bind");
        close(server_fd);
        return;
    }

    if (listen(server_fd, 10) < 0) {
        perror("listen");
        close(server_fd);
        return;
    }

    std::cout << "WebSocket server listening on port " << port << std::endl;

    while (running) {
        sockaddr_in client_addr{};
        socklen_t client_addrlen = sizeof(client_addr);
        int client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &client_addrlen);
        if (client_fd < 0) {
            if (running) perror("accept");
            continue;
        }

        // 处理WebSocket握手
        char buffer[4096];
        ssize_t len = recv(client_fd, buffer, sizeof(buffer) - 1, 0);
        if (len <= 0) {
            close(client_fd);
            continue;
        }
        buffer[len] = '\0';

        if (strstr(buffer, "Upgrade: websocket") != nullptr) {
            char* key_start = strstr(buffer, "Sec-WebSocket-Key:");
            if (key_start) {
                key_start += strlen("Sec-WebSocket-Key:");
                char* key_end = strstr(key_start, "\r\n");
                if (key_end) {
                    std::string key(key_start, key_end - key_start);
                    key.erase(0, key.find_first_not_of(" \t"));
                    key.erase(key.find_last_not_of(" \t") + 1);

                    const std::string magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
                    std::string combined = key + magic;
                    unsigned char sha1[20];
                    SHA1(reinterpret_cast<const unsigned char*>(combined.c_str()), combined.size(), sha1);
                    std::string accept_key = base64_encode(sha1, 20);

                    std::string response = 
                        "HTTP/1.1 101 Switching Protocols\r\n"
                        "Upgrade: websocket\r\n"
                        "Connection: Upgrade\r\n"
                        "Sec-WebSocket-Accept: " + accept_key + "\r\n\r\n";
                    send(client_fd, response.c_str(), response.size(), 0);
                    
                    {
                        std::lock_guard<std::mutex> lock(clients_mutex);
                        clients.insert(client_fd);
                    }
                    continue;
                    if (send(client_fd, response.c_str(), response.size(), 0) > 0) {
            {
                std::lock_guard<std::mutex> lock(clients_mutex);
                clients.insert(client_fd);
            }
            
            // 创建新线程处理客户端
            std::thread client_thread(&WebSocketServer::handleClient, this, client_fd);
            client_thread.detach();
        }
                }
            }
            

        }

        close(client_fd);
    }
}

void WebSocketServer::broadcast(const LogMessage& message) {
    // 将消息转换为JSON
    std::string json = 
        "{\"client_ip\":\"" + message.client_ip + "\","
        "\"method\":\"" + message.method + "\","
        "\"url\":\"" + message.url + "\","
        "\"result\":\"" + message.result + "\","
        "\"port\":" + std::to_string(message.port) + ","
        "\"timestamp\":\"" + message.timestamp + "\"}";
    // WebSocket帧格式: FIN=1, Opcode=1 (text), no mask
    std::vector<unsigned char> frame;
    frame.push_back(0x81); // FIN + text frame
    
    if (json.size() <= 125) {
        frame.push_back(static_cast<unsigned char>(json.size()));
    } else if (json.size() <= 65535) {
        frame.push_back(126);
        frame.push_back((json.size() >> 8) & 0xFF);
        frame.push_back(json.size() & 0xFF);
    } else {
        // 不支持过大的消息
        return;
    }
    
    // 添加数据
    frame.insert(frame.end(), json.begin(), json.end());
    std::cout<<json<<std::endl;
    // 发送给所有客户端
    std::lock_guard<std::mutex> lock(clients_mutex);
    auto it = clients.begin();
    while (it != clients.end()) {
        int client_fd = *it;
        if (send(client_fd, frame.data(), frame.size(), 0) <= 0) {
            close(client_fd);
            it = clients.erase(it);
        } else {
            ++it;
        }
    }
}
WebSocketServer::WebSocketServer(int port) : port(port), running(false) {}

WebSocketServer::~WebSocketServer() {
    stop();
}
void log_processor_thread(WebSocketServer* ws_server) {
    while (log_thread_running) {
        std::unique_lock<std::mutex> lock(queue_mutex);
        queue_cv.wait(lock, []{ return !log_queue.empty() || !log_thread_running; });
        
        if (!log_thread_running) break;
        
        while (!log_queue.empty()) {
            LogMessage msg = log_queue.front();
            log_queue.pop();
            lock.unlock();
            
            // 广播日志
            ws_server->broadcast(msg);
            
            lock.lock();
        }
    }
}